# -*- coding: utf-8 -*-
"""
   File Name:  train.py
   Author :    liccoo
   Time:       2022/8/24 12:48
"""
import argparse
import os

import torch

from utils.load_dataset import forward_dataloader, inverse_dataloader
from utils.models import ForwardModel, InverseModel
from utils.path import logs_exp, extract_model_path
from utils.print_save_hyperparameter import print_save_hyperparas
from utils.train_model import train_model, train_tandem_model

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # 通用超参数
    parser.add_argument('--lr', type=float, default=1e-4, help='The initial learning rate')
    parser.add_argument('--tensorboard', action='store_false', help='Whether tensorboard will start automatically')
    parser.add_argument('--port', type=int, default=6006, help='The port of tensorboard')
    parser.add_argument('--batch-size', type=int, default=64, help='Total batch size for all GPUs')
    parser.add_argument('--epochs', type=int, default=10000, help='The iterations of epoch')
    parser.add_argument('--save-num', type=int, default=5, help='The number of models saved each training')
    parser.add_argument('--device', type=str, default='cuda:0', help='cuda:0, cuda:2, ..., cuda:n, cpu')
    parser.add_argument('--num-workers', type=int, default=2, help='0, 2, 4, 8, 16')
    parser.add_argument('--shuffle', type=bool, default=True, help='Whether the dataset will be shuffled')
    parser.add_argument('--step-size', type=int, default=100, help='Step interval for learning rate decay')
    parser.add_argument('--gamma', type=float, default=0.9, help='Learning rate decay coefficient')
    parser.add_argument('--new-dir', type=str, default='',
                        help='Optional, the additional name of the parent directory path of exp')
    # 定制超参数
    parser.add_argument('--model', type=str, required=True, help='|forward|inverse|tandem|')
    parser.add_argument('--forward-exp-parent', type=str, default='',
                        help='The custom parent directory name of the exp of the forward inverse_model. '
                             'Before setting, make sure that the directory exists')
    parser.add_argument('--exp-order', type=int, default=0,
                        help='The path of the forward inverse_model loaded,default=0 indicates the latest exp.')
    parser.add_argument('--forward-model-order', type=int, default=0,
                        help='The order of the forward_model which are saved. '
                             'default=0 indicates the latest forward_model')
    args = parser.parse_args()

    # 获取当前项目的绝对路径
    project_dir_path = os.path.dirname(os.path.realpath('__file__'))

    # 数据集路径
    train_dataset_path = os.path.join(project_dir_path, 'dataset', 'train_dataset.csv')
    test_dataset_path = os.path.join(project_dir_path, 'dataset', 'test_dataset.csv')

    # 设置存放训练 log 文件的根目录
    logs_exp_path = logs_exp(project_dir_path, args, 'train')

    # 打印训练的模型、是否使用 GPU
    print('-' * 100)  # 打印分割线
    print(f'  Model: {args.model}   Training model using {args.device}')
    print('-' * 100, '\n')  # 打印分割线

    # 打印和保存超参数
    print_save_hyperparas(args, logs_exp_path)

    if args.model == 'forward':  # 前向网络
        # 加载 dataloader
        train_dataloader = forward_dataloader(train_dataset_path, args)
        test_dataloader = forward_dataloader(test_dataset_path, args)
        # 实例化网络模型
        model = ForwardModel()
        # 训练网络
        train_model(args, logs_exp_path, train_dataloader, test_dataloader, model)

    if args.model == 'inverse':  # 逆向网络
        # 加载 dataloader
        train_dataloader = inverse_dataloader(train_dataset_path, args)
        test_dataloader = inverse_dataloader(test_dataset_path, args)
        # 实例化网络模型
        model = InverseModel()
        # 训练网络
        train_model(args, logs_exp_path, train_dataloader, test_dataloader, model)

    elif args.model == 'tandem':  # 串联网络
        # 加载 dataloader
        train_dataloader = inverse_dataloader(train_dataset_path, args)
        test_dataloader = inverse_dataloader(test_dataset_path, args)
        # 读取用于串联的前向网络
        forward_model_path = extract_model_path(project_dir_path, 'forward',
                                                args.exp_order, args.forward_model_order, args.forward_exp_parent)
        # 读取训练好的正向模型
        if args.device == 'cpu':
            forward_model = torch.load(forward_model_path, map_location=torch.device('cpu'))
        else:
            forward_model = torch.load(forward_model_path)
        # 锁定梯度
        for para in forward_model.parameters():
            para.requires_grad = False
        # 实例化网络模型
        inverse_model = InverseModel()
        # 训练网络
        train_tandem_model(args, logs_exp_path, train_dataloader, test_dataloader, forward_model, inverse_model)
